from deeprobust.graph.targeted_attack import IGAttack
from deeprobust.graph.utils import *
import load_data

print("CUDA:%s" % torch.cuda.is_available())
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


def features(name, data_type):
    adj, features, labels, idx_train, idx_val, idx_test, idx_unlabeled = load_data.load(name)
    # 加载代理模型
    surrogate = torch.load('./save_model/ig_' + name + '_model.pth')
    surrogate.load_state_dict(torch.load('./save_model/ig_' + name + '_model_params.pth'))
    degrees = adj.sum(0).A1
    if data_type == 'train':
        node_list = split_list(idx_train.tolist(), 10)
    else:
        node_list = split_list(idx_test.tolist(), 10)
    modified_features = features.tolil()
    modified_adj = adj.tolil()
    for i in range(len(node_list)):
        for target_node in node_list[i]:
            n_perturbations = int(degrees[target_node])
            model = IGAttack(surrogate, nnodes=adj.shape[0], attack_structure=False, attack_features=True,
                             device=str(device)).to(device)
            model.attack(features, adj, labels, idx_train, target_node, n_perturbations, steps=20)
            modified_adj.rows[target_node] = model.modified_adj.tolil().rows[target_node]
            modified_features.rows[target_node] = model.modified_features.tolil().rows[target_node]
        adj_table = './outputs/ig_attack/' + name + '/' + data_type + '/features/' + str(
                (i + 1) * 10) + '%_nodes_modified_adj.pt'
        features_table = './outputs/ig_attack/' + name + '/' + data_type + '/features/' + str(
                (i + 1) * 10) + '%_nodes_modified_features.pt'
        torch.save(modified_adj, adj_table)
        torch.save(modified_features, features_table)
        print(adj_table)
        print(features_table)


def split_list(lst, n):
    # 计算每份应该有多少元素
    avg = len(lst) / float(n)
    # 计算每个元素的下标范围
    out = []
    last = 0.0
    while last < len(lst):
        out.append(lst[int(last):int(last + avg)])
        last += avg
    return out


def structure(name, data_type):
    adj, features, labels, idx_train, idx_val, idx_test, idx_unlabeled = load_data.load(name)
    # 加载代理模型
    surrogate = torch.load('./save_model/ig_' + name + '_model.pth')
    surrogate.load_state_dict(torch.load('./save_model/ig_' + name + '_model_params.pth'))
    degrees = adj.sum(0).A1
    if data_type == 'train':
        node_list = split_list(idx_train.tolist(), 10)
    else:
        node_list = split_list(idx_test.tolist(), 10)
    modified_features = features.tolil()
    modified_adj = adj.tolil()
    for i in range(len(node_list)):
        for target_node in node_list[i]:
            n_perturbations = int(degrees[target_node])
            model = IGAttack(surrogate, nnodes=adj.shape[0], attack_structure=True, attack_features=False,
                             device=str(device)).to(device)
            model.attack(features, adj, labels, idx_train, target_node, n_perturbations, steps=20)
            modified_adj.rows[target_node] = model.modified_adj.tolil().rows[target_node]
            modified_features.rows[target_node] = model.modified_features.tolil().rows[target_node]

        adj_table = './outputs/ig_attack/' + name + '/' + data_type + '/structure/' + str(
            (i + 1) * 10) + '%_nodes_modified_adj.pt'
        features_table = './outputs/ig_attack/' + name + '/' + data_type + '/structure/' + str(
            (i + 1) * 10) + '%_nodes_modified_features.pt'
        torch.save(modified_adj, adj_table)
        torch.save(modified_features, features_table)
        print(adj_table)
        print(features_table)


if __name__ == '__main__':
    datas = ['cora', 'citeseer']
    for data in datas:
        for i in ['train', 'test']:
            features(data, i)
            structure(data, i)
